#coding=utf-8
#---------------------------------------------------
#导入库
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from sklearn.model_selection import cross_val_score, KFold
import kashgari
from kashgari.embeddings import BERTEmbedding
from kashgari.tasks.labeling import BiLSTM_CRF_Model
#----------------数据探索----------------
#数据预处理
#加载训练集
train_df = pd.read_csv('./基于论文摘要的文本分类与查询性问答公开数据/train.csv', sep=',')
#加载测试集
test_df = pd.read_csv('./基于论文摘要的文本分类与查询性问答公开数据/test.csv', sep=',')

#EDA数据探索性分析
train_df.head()

test_df.head()

#----------------特征工程----------------
#将Topic(Label)编码
train_df['Topic(Label)'], lbl = pd.factorize(train_df['Topic(Label)'])

#将论文的标题与摘要组合为 text 特征
train_df['Title'] = train_df['Title'].apply(lambda x: x.strip())
train_df['Abstract'] = train_df['Abstract'].fillna('').apply(lambda x: x.strip())
train_df['text'] = train_df['Title'] + ' ' + train_df['Abstract']
train_df['text'] = train_df['text'].str.lower()

test_df['Title'] = test_df['Title'].apply(lambda x: x.strip())
test_df['Abstract'] = test_df['Abstract'].fillna('').apply(lambda x: x.strip())
test_df['text'] = test_df['Title'] + ' ' + test_df['Abstract']
test_df['text'] = test_df['text'].str.lower()
print(test_df['text'])

#----------------模型训练----------------
kashgari.config.use_cudnn_cell = True

print(type(test_df['text']))
train_x = range(test_df['text'])
train_y = range(train_df['Topic(Label)'])
test_x = range(test_df['text'])

#
# print(train_x)
# print(train_y)
print(type(train_df['Topic(Label)']))

bert_embed = BERTEmbedding('F:\BaiduNetdiskDownload\chinese_bert_wwm_ext_L-12_H-768_A-12',task=kashgari.LABELING)

# 还可以选择 `CNN_LSTM_Model`, `BiLSTM_Model`, `BiGRU_Model` 或 `BiGRU_CRF_Model`
model = BiLSTM_CRF_Model(bert_embed)
model.fit(train_x,
          train_y,
          train_x,
          train_y,
          epochs=1,
          batch_size=512)

model.evaluate(train_x,train_df['Topic(Label)'])
#用于获取每个交叉验证的得分,然后根据得分score来选择合适的超参数,通常需要编写手动完成交叉
cvs = cross_val_score(model, train_x, train_y, cv=5)
print(cvs)

test_df['Topic(Label)'] = model.predict(test_x)

#----------------结果输出----------------
print(test_df['Topic(Label)'])
test_df['Topic(Label)'] = test_df['Topic(Label)'].apply(lambda x: lbl[x])
test_df[['Topic(Label)']].to_csv('submit-tfidf-KFold.csv', index=None)